// The module 'vscode' contains the VS Code extensibility API
// Import the module and reference it with the alias vscode in your code below
import * as vscode from 'vscode';
import * as fs from 'fs';
import * as path from 'path';

// This method is called when your extension is activated
// Your extension is activated the very first time the command is executed
export function activate(context: vscode.ExtensionContext) {

	// Use the console to output diagnostic information (console.log) and errors (console.error)
	// This line of code will only be executed once when your extension is activated
	console.log('Congratulations, your extension "elf64-reader" is now active!');

	// The command has been defined in the package.json file
	// Now provide the implementation of the command with registerCommand
	// The commandId parameter must match the command field in package.json
	const disposable = vscode.commands.registerCommand('elf64-reader.helloWorld', () => {
		// The code you place here will be executed every time your command is executed
		// Display a message box to the user
		vscode.window.showInformationMessage('Hello World from elf64-reader!');
	});

	context.subscriptions.push(disposable);

	// Register a new command to read ELF64 file
	const readElfCommand = vscode.commands.registerCommand('elf64-reader.readElfFile', async () => {
		const fileUri = await vscode.window.showOpenDialog({
			canSelectFolders: false,
			canSelectFiles: true,
			canSelectMany: false,
			filters: {
				'ELF64 Files': ['o','elf']
			}
		});

		if (fileUri && fileUri.length > 0) {
			readElfFile(fileUri[0].fsPath);
		}
	});

	context.subscriptions.push(readElfCommand);
}

// This method is called when your extension is deactivated
export function deactivate() {}

// Function to read and parse ELF64 file
function readElfFile(filePath: string) {
    try {
        const fileBuffer = fs.readFileSync(filePath);
	const dataView = new DataView(fileBuffer.buffer, fileBuffer.byteOffset, fileBuffer.byteLength);

        // Parse ELF Header
	const elfHeader = parseElfHeader(dataView);
        if (!elfHeader) {
            throw new Error('Failed to parse ELF header');
        }
	console.log('ELF Header parsed successfully');

        // Parse Section Headers
	const sectionHeaders = parseSectionHeaders(dataView, elfHeader);
	console.log(`Parsed ${sectionHeaders.length} section headers`);

        // Parse Program Headers
	const programHeaders = parseProgramHeaders(dataView, elfHeader);
        if (!programHeaders) {
            throw new Error('Failed to parse program headers');
        }
	console.log(`Parsed ${programHeaders.length} program headers`);

        // Parse Symbol Table
	const symbolTable = parseSymbolTable(dataView, elfHeader, sectionHeaders);
        console.log(`Parsed ${symbolTable.length} symbols`);
	// Show parsed information in a webview

	// Show parsed information in a webview
	// Precompute section names for webview
    const shstrtabSection = sectionHeaders[elfHeader.eShstrndx];
    const shstrtab = parseStringTable(dataView, shstrtabSection.shOffset, shstrtabSection.shSize);
    const namedSectionHeaders = sectionHeaders.map(sh => ({
        ...sh,
        name: shstrtab.get(sh.shName) || ''
    }));

    // Precompute symbol names and section names
    // Find symbol table section (.symtab)
    // Find symbol table section by type first, then by name
    const symtabByType = sectionHeaders.find(sh => sh.shType === 2) || sectionHeaders.find(sh => sh.shType === 11);
    const symtabByName = namedSectionHeaders.find(sh => sh.name === '.symtab') || namedSectionHeaders.find(sh => sh.name === '.dynsym');
    const symtabSection = symtabByType ? namedSectionHeaders[sectionHeaders.indexOf(symtabByType)] : symtabByName;

    // Get string table section from symtab's shLink with fallbacks
    let strtabSection = symtabSection ? namedSectionHeaders[symtabSection.shLink] : null;

    // If linked section isn't a string table, try to find by type or name
    if (strtabSection && strtabSection.shType !== 3) {
        strtabSection = namedSectionHeaders.find(sh => sh.shType === 3);
    }
    if (!strtabSection) {
        strtabSection = namedSectionHeaders.find(sh => sh.name === '.strtab') || 
                        namedSectionHeaders.find(sh => sh.name === '.dynstr');
    }

    const strtab = strtabSection ? parseStringTable(dataView, strtabSection.shOffset, strtabSection.shSize) : new Map();
    const namedSymbolTable = symbolTable.map(sym => ({
        ...sym,
        name: strtab.get(sym.stName) || ''
    }));

    showElfInfoInWebview({ elfHeader, sectionHeaders: namedSectionHeaders, programHeaders, symbolTable: namedSymbolTable });
    } catch (error) {
        console.error('Error reading ELF file:', error);
        vscode.window.showErrorMessage(`Failed to read ELF file: ${error instanceof Error ? error.message : String(error)}`);
    }
}

// Function to parse ELF Header
function parseElfHeader(dataView: DataView) {
    const elfHeader: any = {};

    // ELF Identification
    elfHeader.eiMagic = dataView.getUint32(0, true);
    elfHeader.eiClass = dataView.getUint8(4); // 1 for 32-bit, 2 for 64-bit
    elfHeader.eiData = dataView.getUint8(5); // 1 for little-endian, 2 for big-endian
    const littleEndian = elfHeader.eiData === 1;
    elfHeader.eiVersion = dataView.getUint8(6);
    elfHeader.eiOsAbi = dataView.getUint8(7);
    elfHeader.eiAbiVersion = dataView.getUint8(8);

    // ELF Header fields
    elfHeader.eType = dataView.getUint16(16, littleEndian);
    elfHeader.eMachine = dataView.getUint16(18, littleEndian);
    elfHeader.eVersion = dataView.getUint32(20, littleEndian);
    elfHeader.eEntry = dataView.getBigUint64(24, littleEndian);
    elfHeader.ePhoff = dataView.getBigUint64(32, littleEndian);
    elfHeader.eShoff = dataView.getBigUint64(40, littleEndian);
    elfHeader.eFlags = dataView.getUint32(48, littleEndian);
    elfHeader.eEhsize = dataView.getUint16(52, littleEndian);
    elfHeader.ePhentsize = dataView.getUint16(54, littleEndian);
    elfHeader.ePhnum = dataView.getUint16(56, littleEndian);
    elfHeader.eShentsize = dataView.getUint16(58, littleEndian);
    elfHeader.eShnum = dataView.getUint16(60, littleEndian);
    elfHeader.eShstrndx = dataView.getUint16(62, littleEndian);

    return elfHeader;
}

// Function to parse Section Headers
function parseSectionHeaders(dataView: DataView, elfHeader: any) {
    const sectionHeaders: any[] = [];
    const sectionHeaderOffset = elfHeader.eShoff;
    const sectionHeaderSize = elfHeader.eShentsize;
    const sectionHeaderCount = elfHeader.eShnum;

    for (let i = 0; i < sectionHeaderCount; i++) {
        const offset = sectionHeaderOffset + BigInt(i) * BigInt(sectionHeaderSize);
        const sectionHeader: any = {};

        const littleEndian = elfHeader.eiData === 1;
        sectionHeader.shName = dataView.getUint32(Number(offset), littleEndian);
        sectionHeader.shType = dataView.getUint32(Number(offset + 4n), littleEndian);
        sectionHeader.shFlags = dataView.getBigUint64(Number(offset + 8n), littleEndian);
        sectionHeader.shAddr = dataView.getBigUint64(Number(offset + 16n), littleEndian);
        sectionHeader.shOffset = dataView.getBigUint64(Number(offset + 24n), littleEndian);
        sectionHeader.shSize = dataView.getBigUint64(Number(offset + 32n), littleEndian);
        sectionHeader.shLink = dataView.getUint32(Number(offset + 40n), littleEndian);
        sectionHeader.shInfo = dataView.getUint32(Number(offset + 44n), littleEndian);
        sectionHeader.shAddralign = dataView.getBigUint64(Number(offset + 48n), littleEndian);
        sectionHeader.shEntsize = dataView.getBigUint64(Number(offset + 56n), littleEndian);

        sectionHeaders.push(sectionHeader);
    }

    return sectionHeaders;
}

// Function to parse Program Headers
function parseProgramHeaders(dataView: DataView, elfHeader: any) {
    const programHeaders: any[] = [];
    const programHeaderOffset = elfHeader.ePhoff;
    const programHeaderSize = elfHeader.ePhentsize;
    const programHeaderCount = elfHeader.ePhnum;

    for (let i = 0; i < programHeaderCount; i++) {
        const offset = programHeaderOffset + BigInt(i) * BigInt(programHeaderSize);
        const programHeader: any = {};

        const littleEndian = elfHeader.eiData === 1;
        programHeader.pType = dataView.getUint32(Number(offset), littleEndian);
        programHeader.pFlags = dataView.getUint32(Number(offset + 4n), littleEndian);
        programHeader.pOffset = dataView.getBigUint64(Number(offset + 8n), littleEndian);
        programHeader.pVaddr = dataView.getBigUint64(Number(offset + 16n), littleEndian);
        programHeader.pPaddr = dataView.getBigUint64(Number(offset + 24n), littleEndian);
        programHeader.pFilesz = dataView.getBigUint64(Number(offset + 32n), littleEndian);
        programHeader.pMemsz = dataView.getBigUint64(Number(offset + 40n), littleEndian);
        programHeader.pAlign = dataView.getBigUint64(Number(offset + 48n), littleEndian);

        programHeaders.push(programHeader);
    }

    return programHeaders;
}

// Helper function to parse string table
function parseStringTable(dataView: DataView, offset: bigint, size?: bigint): Map<number, string> {
    const stringTable = new Map<number, string>();
    let currentOffset = Number(offset);
    let stringStart = currentOffset;

    const endOffset = size ? Number(offset) + Number(size) : dataView.byteLength;
    while (currentOffset < endOffset) {
        const charCode = dataView.getUint8(currentOffset);
        if (charCode === 0) {
            if (currentOffset > stringStart) {
                const str = new TextDecoder().decode(
                    new Uint8Array(dataView.buffer, dataView.byteOffset + stringStart, currentOffset - stringStart)
                );
                stringTable.set(stringStart - Number(offset), str);
            }
            stringStart = currentOffset + 1;
        }
        currentOffset++;
    }
    return stringTable;
}

// Mapping tables for human-readable values
const SECTION_TYPE_MAP: Record<number, string> = {
    0: 'SHT_NULL',
    1: 'SHT_PROGBITS',
    2: 'SHT_SYMTAB',
    3: 'SHT_STRTAB',
    4: 'SHT_RELA',
    5: 'SHT_HASH',
    6: 'SHT_DYNAMIC',
    7: 'SHT_NOTE',
    8: 'SHT_NOBITS',
    9: 'SHT_REL',
    10: 'SHT_SHLIB',
    11: 'SHT_DYNSYM',
    14: 'SHT_INIT_ARRAY',
    15: 'SHT_FINI_ARRAY',
    16: 'SHT_PREINIT_ARRAY',
    17: 'SHT_GROUP',
    18: 'SHT_SYMTAB_SHNDX',
    0x70000000: 'SHT_LOPROC',
    0x7fffffff: 'SHT_HIPROC',
    0x80000000: 'SHT_LOUSER',
    0x8fffffff: 'SHT_HIUSER'
};

const SECTION_FLAG_MAP: Record<number, string> = {
    0: ' ',
    1: 'W (write)',
    2: 'A (alloc)',
    4: 'X (execute)',
    8: 'M (merge)',
    0x10: 'S (strings)',
    0x20: 'I (info link)',
    0x40: 'L (link order)',
    0x80: 'O (extra OS processing required)',
    0x100: 'G (group)',
    0x200: 'T (TLS)',
    0x400: 'C (compressed)',
    0x800: 'X (exclude)',
    0x1000: 'S (OS specific)',
    0x2000: 'E (OS specific)',
    0x4000: 'D (OS specific)',
    0x8000: 'L (OS specific)'
};

const PROGRAM_TYPE_MAP: Record<number, string> = {
    0: 'PT_NULL',
    1: 'PT_LOAD',
    2: 'PT_DYNAMIC',
    3: 'PT_INTERP',
    4: 'PT_NOTE',
    5: 'PT_SHLIB',
    6: 'PT_PHDR',
    7: 'PT_TLS',
    0x60000000: 'PT_LOOS',
    0x6fffffff: 'PT_HIOS',
    0x70000000: 'PT_LOPROC',
    0x7fffffff: 'PT_HIPROC'
};

const PROGRAM_FLAG_MAP: Record<number, string> = {
    0: ' ',
    1: 'R (read)',
    2: 'W (write)',
    4: 'E (execute)'
};

const SYMBOL_TYPE_MAP: Record<number, string> = {
    0: 'NOTYPE',
    1: 'OBJECT',
    2: 'FUNC',
    3: 'SECTION',
    4: 'FILE',
    5: 'COMMON',
    6: 'TLS'
};

const SYMBOL_BINDING_MAP: Record<number, string> = {
    0: 'LOCAL',
    1: 'GLOBAL',
    2: 'WEAK',
    3: 'LOPROC',
    4: 'HIPROC'
};

const ELF_TYPE_MAP: Record<number, string> = {
    0: 'ET_NONE',
    1: 'ET_REL',
    2: 'ET_EXEC',
    3: 'ET_DYN',
    4: 'ET_CORE',
    0xff00: 'ET_LOOS',
    0xffff: 'ET_HIOS',
    0x10000: 'ET_LOPROC',
    0x1ffff: 'ET_HIPROC'
};

const ELF_MACHINE_MAP: Record<number, string> = {
    0: 'EM_NONE',
    3: 'EM_386',
    8: 'EM_MIPS',
    40: 'EM_ARM',
    62: 'EM_X86_64',
    183: 'EM_AARCH64'
};

// Function to parse Symbol Table
function parseSymbolTable(dataView: DataView, elfHeader: any, sectionHeaders: any[]) {
    const symbolTable: any[] = [];

    // Get section header string table

    const shstrtabSection = sectionHeaders[elfHeader.eShstrndx];
    const shstrtab = parseStringTable(dataView, shstrtabSection.shOffset);

    // Find symbol table section (.symtab)
    const symtabSection = sectionHeaders.find(sh => sh.shType === 2); // SHT_SYMTAB = 2
    if (!symtabSection) {
        console.warn('Symbol table section (.symtab) not found');
        return [];
    }

    // Get symbol string table
    const strtabSection = sectionHeaders[symtabSection.shLink];
    const strtab = parseStringTable(dataView, strtabSection.shOffset);

    const symbolTableOffset = symtabSection.shOffset;
    const symbolTableEntrySize = BigInt(symtabSection.shEntsize);
    const symbolTableCount = symtabSection.shSize / symbolTableEntrySize;

    let i = 0n;
    while (i < symbolTableCount) {
        const offset = symbolTableOffset + i * symbolTableEntrySize;
        const symbol: any = {};

        const littleEndian = elfHeader.eiData === 1;
        const nameOffset = dataView.getUint32(Number(offset), littleEndian);
        symbol.stName = nameOffset;
        symbol.stValue = dataView.getBigUint64(Number(offset + 8n), littleEndian);
        symbol.stSize = dataView.getBigUint64(Number(offset + 16n), littleEndian);
        symbol.stInfo = dataView.getUint8(Number(offset + 24n));
        symbol.stOther = dataView.getUint8(Number(offset + 25n));
        symbol.stShndx = dataView.getUint16(Number(offset + 26n), littleEndian);

        symbolTable.push(symbol);
        i++;
    }

    return symbolTable;
}


// Function to show parsed ELF information in a webview
function showElfInfoInWebview(elfInfo: any & { dataView: DataView }) {
	const panel = vscode.window.createWebviewPanel(
		'elfInfo',
		'ELF64 File Information',
		vscode.ViewColumn.One,
		{}
	);

	// Generate HTML content to display ELF information
	const htmlContent = generateHtmlContent(elfInfo);
	panel.webview.html = htmlContent;
}

// Function to generate HTML content for displaying ELF information
function generateHtmlContent(elfInfo: any): string {
	// Generate HTML content with tables for ELF Header, Section Headers, Program Headers, and Symbol Table
	return `
		<!DOCTYPE html>
		<html lang="en">
		<head>
			<meta charset="UTF-8">
			<meta name="viewport" content="width=device-width, initial-scale=1.0">
			<title>ELF64 File Information</title>
		</head>
		<body>
			<h1>ELF Header</h1>
			<table border="1">
				<tr><th>Field</th><th>Value</th></tr>
				<tr><td>Magic</td><td>7f 45 4c 46 (ELF)</td></tr>
				<tr><td>Class</td><td>${elfInfo.elfHeader.eiClass === 1 ? 'ELF32' : 'ELF64'}</td></tr>
				<tr><td>Data</td><td>${elfInfo.elfHeader.eiData === 1 ? '2\'s complement, little endian' : '2\'s complement, big endian'}</td></tr>
				<tr><td>Version</td><td>${elfInfo.elfHeader.eiVersion === 1 ? '1 (current)' : elfInfo.elfHeader.eiVersion}</td></tr>
				<tr><td>OS/ABI</td><td>${elfInfo.elfHeader.eiOsAbi === 0 ? 'UNIX - System V' : elfInfo.elfHeader.eiOsAbi}</td></tr>
				<tr><td>ABI Version</td><td>${elfInfo.elfHeader.eiAbiVersion}</td></tr>
				<tr><td>Type</td><td>${ELF_TYPE_MAP[elfInfo.elfHeader.eType] || elfInfo.elfHeader.eType}</td></tr>
				<tr><td>Machine</td><td>${ELF_MACHINE_MAP[elfInfo.elfHeader.eMachine] || elfInfo.elfHeader.eMachine}</td></tr>
				<tr><td>Version</td><td>${elfInfo.elfHeader.eVersion === 1 ? '0x1 (current)' : elfInfo.elfHeader.eVersion.toString(16)}</td></tr>
				<tr><td>Entry Point</td><td>0x${elfInfo.elfHeader.eEntry.toString(16)}</td></tr>
				<tr><td>Program Header Offset</td><td>${elfInfo.elfHeader.ePhoff} (bytes into file)</td></tr>
				<tr><td>Section Header Offset</td><td>${elfInfo.elfHeader.eShoff} (bytes into file)</td></tr>
				<tr><td>Flags</td><td>0x${elfInfo.elfHeader.eFlags.toString(16)}</td></tr>
				<tr><td>ELF Header Size</td><td>${elfInfo.elfHeader.eEhsize} (bytes)</td></tr>
				<tr><td>Program Header Entry Size</td><td>${elfInfo.elfHeader.ePhentsize} (bytes)</td></tr>
				<tr><td>Number of Program Headers</td><td>${elfInfo.elfHeader.ePhnum}</td></tr>
				<tr><td>Section Header Entry Size</td><td>${elfInfo.elfHeader.eShentsize} (bytes)</td></tr>
				<tr><td>Number of Section Headers</td><td>${elfInfo.elfHeader.eShnum}</td></tr>
				<tr><td>Section Name String Table Index</td><td>${elfInfo.elfHeader.eShstrndx}</td></tr>
			</table>

			<h1>Section Headers</h1>
			<table border="1">
				<tr><th>Index</th><th>Name</th><th>Type</th><th>Flags</th><th>Address</th><th>Offset</th><th>Size</th><th>Link</th><th>Info</th><th>Address Alignment</th><th>Entry Size</th></tr>
				${elfInfo.sectionHeaders.map((sh: any, index: number) => {
						// Convert flags to human-readable string
						let flags = '';
						if (sh.shFlags & 1n) {
							flags += 'W';
						}
						if (sh.shFlags & 2n) {
							flags += 'A';
						}
						if (sh.shFlags & 4n) {
							flags += 'X';
						}
						if (flags === '') {
							flags = ' ';
						}

						return `
					<tr>
						<td>${index}</td>
						<td>${sh.name}</td>
						<td>${SECTION_TYPE_MAP[sh.shType] || `0x${sh.shType.toString(16)}`}</td>
						<td>[${flags}] ${Array.from(Object.entries(SECTION_FLAG_MAP))
							.filter(([flag, desc]) => BigInt(flag) & sh.shFlags && desc !== ' ')
							.map(([_, desc]) => desc)
							.join(', ')}</td>
						<td>0x${sh.shAddr.toString(16)}</td>
						<td>${sh.shOffset} (bytes into file)</td>
						<td>${sh.shSize.toString()} (bytes)</td>
						<td>${sh.shLink} (${sh.shLink === 0 ? 'UND' : sh.shLink === 0xfff1 ? 'LOPROC' : sh.shLink === 0xfff2 ? 'HIPROC' : sh.shLink === 0xfff3 ? 'LOOS' : sh.shLink === 0xffff ? 'HIOS' : elfInfo.sectionHeaders[sh.shLink]?.name || ''})</td>
						<td>${sh.shInfo}</td>
						<td>${sh.shAddralign.toString()}</td>
						<td>${sh.shEntsize.toString()}</td>
					</tr>
				`;}).join('')})()}
			</table>

			<h1>Program Headers</h1>
			<table border="1">
				<tr><th>Index</th><th>Type</th><th>Flags</th><th>Offset</th><th>Virtual Address</th><th>Physical Address</th><th>File Size</th><th>Memory Size</th><th>Alignment</th></tr>
				${elfInfo.programHeaders.map((ph: any, index: number) => {
					// Convert flags to human-readable string
					let flagsStr = '';
					if (ph.pFlags & 1) { flagsStr += 'R'; }
					if (ph.pFlags & 2) { flagsStr += 'W'; }
					if (ph.pFlags & 4) { flagsStr += 'E'; }
					if (flagsStr === '') { flagsStr = ' '; }

					return `
					<tr>
						<td>${index}</td>
						<td>${PROGRAM_TYPE_MAP[ph.pType] || `0x${ph.pType.toString(16)}`}</td>
						<td>[${flagsStr}] ${Array.from(Object.entries(PROGRAM_FLAG_MAP))
							.filter(([flag, desc]) => Number(flag) & ph.pFlags && desc !== ' ')
							.map(([_, desc]) => desc)
							.join(', ')}</td>
						<td>${ph.pOffset} (bytes into file)</td>
						<td>0x${ph.pVaddr.toString(16)}</td>
						<td>0x${ph.pPaddr.toString(16)}</td>
						<td>${ph.pFilesz.toString()} (bytes)</td>
						<td>${ph.pMemsz.toString()} (bytes)</td>
						<td>${ph.pAlign === 0n ? 1 : ph.pAlign.toString()}</td>
					</tr>
				`;}).join('')}
			</table>

			<h1>Symbol Table</h1>
			<table border="1">
				<tr><th>Index</th><th>Name</th><th>Value</th><th>Size</th><th>Type</th><th>Binding</th><th>Visibility</th><th>Section Index</th></tr>
				${elfInfo.symbolTable.map((sym: any, index: number) => {
					const type = sym.stInfo & 0x0F;
					const binding = (sym.stInfo >> 4) & 0x0F;
					const visibility = sym.stOther & 0x03;
					const visibilityStr = visibility === 0 ? 'default' : visibility === 1 ? 'internal' : visibility === 2 ? 'hidden' : 'protected';

					return `
					<tr>
						<td>${index}</td>
						<td>${sym.name || ' '}</td>
						<td>0x${sym.stValue.toString(16)}</td>
						<td>${sym.stSize.toString()}</td>
						<td>${SYMBOL_TYPE_MAP[type] || `0x${type.toString(16)}`}</td>
						<td>${SYMBOL_BINDING_MAP[binding] || `0x${binding.toString(16)}`}</td>
						<td>${visibilityStr}</td>
						<td>${sym.stShndx === 0 ? 'UND' : sym.stShndx === 0xfff1 ? 'LOPROC' : sym.stShndx === 0xfff2 ? 'HIPROC' : sym.stShndx === 0xfff3 ? 'LOOS' : sym.stShndx === 0xffff ? 'HIOS' : sym.stShndx}</td>
					</tr>
				`;}).join('')}
			</table>
		</body>
		</html>
	`;
}